#include "trace_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
    // 通用计算对齐长度
    // static inline uint64_t AlignUp(uint64_t x, uint64_t a) { return (x + a - 1) / a * a; }

    const uint32_t blockSize = 32;
    const uint32_t coreNum = 8;
    const uint32_t bufferNum = 2;

    static ge::graphStatus TilingFunc(gert::TilingContext* context)
    {
        TraceCustomTilingData tiling;
        // 从输入tensor获取实际的矩阵参数
        auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
        const auto xShape = context->GetInputTensor(0)->GetOriginShape();
        uint32_t matrix_num = xShape.GetDim(0);     // batch size
        uint32_t row_length = xShape.GetDim(1);     // 矩阵行数 
        uint32_t column_length = xShape.GetDim(2);  // 矩阵列数
        // 设置基础参数(仅ub64位)
        uint32_t type_size = sizeof(uint16_t);
        uint32_t align_num = blockSize / type_size;
        uint64_t ub_size = 192 * 1024; 
        uint64_t ub_max = static_cast<uint32_t>(ub_size / type_size / bufferNum);          // 单个核的ub所能处理的最大元素数量
        // 计算对齐填充
        uint32_t align_padding = (align_num - (row_length % align_num)) % align_num;   
        uint32_t align_length = row_length + align_padding;
        // 主核尾核处理元素个数（长度）--以n = 777，half为例，向上取整为784，每个主核处理98个元素
        uint32_t full_tile_length = align_length / coreNum;
        uint32_t last_tile_length = row_length - full_tile_length * (coreNum - 1);
        // 处理特殊情况：如果只有一个核，则主核大小设为0，只用尾核
        if (row_length <= full_tile_length){ 
            full_tile_length = 0;
            last_tile_length = row_length; 
        }
        context->SetBlockDim(coreNum);
        tiling.set_ubMax(ub_max);
        tiling.set_matrixOrder(0); // 按行优先存储
        tiling.set_matrixNum(matrix_num);
        tiling.set_rowLength(row_length);
        tiling.set_columnLength(column_length);
        tiling.set_fullTileLength(full_tile_length);
        tiling.set_lastTileLength(last_tile_length);
        tiling.set_alignNum(align_num);
        tiling.set_alignPadding(align_padding);
        tiling.set_alignLength(align_length);
        tiling.set_typeSize(type_size);
        tiling.set_loopCnt(1);  // 每个核循环次数

        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        currentWorkspace[0] = 0;
        return ge::GRAPH_SUCCESS;
    }
}

namespace ge {
    static ge::graphStatus InferShape(gert::InferShapeContext* context)
    {
        const gert::Shape* x_shape = context->GetInputShape(0);
        gert::Shape* y_shape = context->GetOutputShape(0);
        y_shape->SetDim(0, x_shape->GetDim(0));  
        y_shape->SetDimNum(1);  
        return GRAPH_SUCCESS;
    }
    static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
    {
        const auto inputDataType = context->GetInputDataType(0);
        context->SetOutputDataType(0, inputDataType);
        return ge::GRAPH_SUCCESS;
    }
}

namespace ops {
class TraceCustom : public OpDef {
public:
    explicit TraceCustom(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape).SetInferDataType(ge::InferDataType);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(TraceCustom);
}
